# ==============================================================================
# Docker镜像导出脚本 - 完整版 (PowerShell版本)
# 包含所有docker-compose.yml和save_docker_images.sh中定义的镜像
# ==============================================================================

# 创建输出目录
$OutputDir = "docker_images_backup"
if (-not (Test-Path $OutputDir)) {
    New-Item -ItemType Directory -Path $OutputDir -Force | Out-Null
}

# 定义输出文件名（包含时间戳）
$Timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
$OutputFile = "$OutputDir\yuxi_know_docker_images_$Timestamp.tar"
$ListFile = "$OutputDir\image_list_$Timestamp.txt"
$LoadScript = "$OutputDir\load_images.ps1"

Write-Host "==============================================="
Write-Host "开始导出 Yuxi-Know 项目所有 Docker 镜像" -ForegroundColor Green
Write-Host "==============================================="
Write-Host "输出文件: $OutputFile"
Write-Host ""

# 从 docker-compose.yml 提取的完整镜像列表
$ComposeImages = @(
    # 自定义构建镜像
    "yuxi-api:0.2.0.dev",
    "yuxi-web:0.2.0.dev",
    "mineru-sglang:latest",
    "paddlex:latest",

    # 数据库相关镜像
    "neo4j:5.26",
    "shangor/postgres-for-rag:v1.0",

    # Milvus 相关镜像
    "quay.io/coreos/etcd:v3.5.5",
    "minio/minio:RELEASE.2023-03-20T20-16-18Z",
    "milvusdb/milvus:v2.5.6"
)

# 从 save_docker_images.sh 提取的基础镜像列表
$BaseImages = @(
    "python:3.11-slim",
    "ghcr.io/astral-sh/uv:0.7.2",
    "node:20-alpine",
    "nginx:alpine",
    "neo4j:5.26",
    "quay.io/coreos/etcd:v3.5.5",
    "minio/minio:RELEASE.2023-03-20T20-16-18Z",
    "milvusdb/milvus:v2.5.6"
)

# 合并所有镜像并去重
$AllImages = @(
    # 自定义构建镜像
    "yuxi-api:0.2.0.dev",
    "yuxi-web:0.2.0.dev",
    "mineru-sglang:latest",
    "paddlex:latest",

    # 数据库相关镜像
    "neo4j:5.26",
    "shangor/postgres-for-rag:v1.0",

    # Milvus 相关镜像
    "quay.io/coreos/etcd:v3.5.5",
    "minio/minio:RELEASE.2023-03-20T20-16-18Z",
    "milvusdb/milvus:v2.5.6",

    # 基础镜像（从Dockerfile分析得出）
    "python:3.11-slim",
    "ghcr.io/astral-sh/uv:0.7.2",
    "node:20-alpine",
    "nginx:alpine",

    # 构建依赖的基础镜像
    "lmsysorg/sglang:v0.4.9.post3-cu126",  # mineru.Dockerfile的基础镜像
    "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlex/paddlex:paddlex3.0.1-paddlepaddle3.0.0-gpu-cuda11.8-cudnn8.9-trt8.6"  # paddlex.Dockerfile的基础镜像
)

# 显示将要导出的镜像列表
Write-Host "将要导出的镜像列表："
Write-Host "-----------------------------------------------"
for ($i = 0; $i -lt $AllImages.Count; $i++) {
    Write-Host "$($i+1). $($AllImages[$i])"
}
Write-Host "总共 $($AllImages.Count) 个镜像"
Write-Host ""

# 确认是否继续
$Confirm = Read-Host "确认导出这些镜像吗？(y/N)"
if ($Confirm -notmatch '^[Yy]') {
    Write-Host "导出已取消"
    exit 0
}

# 检查 Docker 是否运行
try {
    $null = docker info 2>$null
} catch {
    Write-Host "错误: Docker 未运行或权限不足" -ForegroundColor Red
    exit 1
}

# 检查现有镜像并准备导出列表
Write-Host "检查本地镜像..."
Write-Host "-----------------------------------------------"
$SuccessImages = @()
$MissingImages = @()

foreach ($Image in $AllImages) {
    # 检查镜像是否已存在 - 使用更简单的匹配方法
    $AllImageTags = docker images --format "{{.Repository}}:{{.Tag}}"
    $Found = $false

    foreach ($Tag in $AllImageTags) {
        if ($Tag -eq $Image) {
            $Found = $true
            break
        }
    }

    if ($Found) {
        Write-Host "✓ 找到镜像: $Image" -ForegroundColor Green
        $SuccessImages += $Image
    } else {
        Write-Host "❌ 缺少镜像: $Image" -ForegroundColor Red
        $MissingImages += $Image
    }
}

Write-Host ""
Write-Host "镜像统计:"
Write-Host "  ✓ 可用镜像: $($SuccessImages.Count) 个" -ForegroundColor Green
Write-Host "  ❌ 缺少镜像: $($MissingImages.Count) 个" -ForegroundColor Red

if ($SuccessImages.Count -eq 0) {
    Write-Host ""
    Write-Host "错误: 没有找到任何可用的本地镜像" -ForegroundColor Red
    Write-Host "建议: 请先运行 'docker compose build' 构建项目镜像" -ForegroundColor Yellow
    exit 1
}

if ($MissingImages.Count -gt 0) {
    Write-Host ""
    Write-Host "缺少的镜像列表:" -ForegroundColor Yellow
    foreach ($Image in $MissingImages) {
        Write-Host "  - $Image" -ForegroundColor White
    }
    Write-Host ""
    Write-Host "注意: 这些镜像将被跳过，如需包含请先拉取或构建" -ForegroundColor Yellow
}

# 确认是否继续
Write-Host ""
$Confirm = Read-Host "是否继续导出找到的 $($SuccessImages.Count) 个镜像？(y/N)"
if ($Confirm -notmatch '^[Yy]') {
    Write-Host "导出已取消"
    exit 0
}

Write-Host ""
Write-Host "开始保存镜像到 tar 文件..."
Write-Host "-----------------------------------------------"

Write-Host "正在保存 $($SuccessImages.Count) 个镜像到 $OutputFile..."

try {
    # 构建docker save命令
    $SaveArgs = $SuccessImages -join " "
    $Command = "docker save $SaveArgs -o `"$OutputFile`""

    Invoke-Expression $Command

    if ($LASTEXITCODE -eq 0) {
        Write-Host "✓ 镜像保存成功" -ForegroundColor Green
    } else {
        Write-Host "❌ 镜像保存失败" -ForegroundColor Red
        exit 1
    }
} catch {
    Write-Host "❌ 镜像保存失败: $($_.Exception.Message)" -ForegroundColor Red
    exit 1
}

# 计算文件大小
if (Test-Path $OutputFile) {
    $FileSize = (Get-Item $OutputFile).Length
    $FileSizeFormatted = [math]::Round($FileSize / 1GB, 2)
    $SizeDisplay = "$FileSizeFormatted GB"
} else {
    $SizeDisplay = "未知"
}

# 生成镜像列表文件
Write-Host "生成镜像列表文件: $ListFile"
@"
Yuxi-Know Docker 镜像导出列表
导出时间: $(Get-Date)
镜像文件: $OutputFile
文件大小: $SizeDisplay

包含的镜像:
"@ | Out-File -FilePath $ListFile -Encoding UTF8

foreach ($Image in $SuccessImages) {
    "  - $Image" | Out-File -FilePath $ListFile -Append -Encoding UTF8
}

# 生成加载脚本
Write-Host "生成加载脚本: $LoadScript"
@"
# Docker 镜像加载脚本 (PowerShell版本)
# 使用方法: .\load_images.ps1 [image_file.tar]

param(
    [string]`$ImageFile = `"`"
)

# 默认镜像文件
if (`$ImageFile -eq `"`") {
    # 查找最新的 tar 文件
    `$ImageFile = Get-ChildItem -Path "docker_images_backup\*.tar" -ErrorAction SilentlyContinue | Sort-Object LastWriteTime -Descending | Select-Object -First 1 -ExpandProperty FullName
}

if (-not `$ImageFile -or -not (Test-Path `$ImageFile)) {
    Write-Host "错误: 找不到镜像文件" -ForegroundColor Red
    Write-Host "使用方法: .\load_images.ps1 [image_file.tar]"
    exit 1
}

Write-Host "正在从 `$ImageFile 加载 Docker 镜像..." -ForegroundColor Green
Write-Host ""

try {
    docker load -i "`$ImageFile"

    if (`$LASTEXITCODE -eq 0) {
        Write-Host ""
        Write-Host "✓ 镜像加载成功!" -ForegroundColor Green
        Write-Host ""
        Write-Host "加载的镜像列表:" -ForegroundColor Cyan
        docker images --format "table {{.Repository}}:{{.Tag}}`t{{.Size}}`t{{.CreatedAt}}" | Select-Object -First 20
    } else {
        Write-Host "❌ 镜像加载失败" -ForegroundColor Red
        exit 1
    }
} catch {
    Write-Host "❌ 镜像加载失败: `$(`$_.Exception.Message)" -ForegroundColor Red
    exit 1
}
"@ | Out-File -FilePath $LoadScript -Encoding UTF8

# 最终输出
Write-Host ""
Write-Host "==============================================="
Write-Host "导出完成！" -ForegroundColor Green
Write-Host "==============================================="
Write-Host "✓ 镜像文件: $OutputFile" -ForegroundColor Green
Write-Host "✓ 文件大小: $SizeDisplay" -ForegroundColor Green
Write-Host "✓ 镜像数量: $($SuccessImages.Count)/$($AllImages.Count)" -ForegroundColor Green
Write-Host "✓ 镜像列表: $ListFile" -ForegroundColor Green
Write-Host "✓ 加载脚本: $LoadScript" -ForegroundColor Green
Write-Host ""
Write-Host "使用方法:" -ForegroundColor Cyan
Write-Host "  1. 加载镜像: .\load_images.ps1" -ForegroundColor White
Write-Host "  2. 或者直接: docker load -i `"$OutputFile`"" -ForegroundColor White
Write-Host ""
Write-Host "建议:" -ForegroundColor Cyan
Write-Host "  - 将 $OutputDir 目录复制到目标机器" -ForegroundColor White
Write-Host "  - 在目标机器上运行加载脚本" -ForegroundColor White
Write-Host "==============================================="

# 显示跳过的镜像（如果有）
$SkippedCount = $AllImages.Count - $SuccessImages.Count
if ($SkippedCount -gt 0) {
    Write-Host ""
    Write-Host "注意: 有 $SkippedCount 个镜像未能拉取，未包含在导出文件中" -ForegroundColor Yellow
}